import numpy as np
import rospy
from sensor_msgs.msg import Image
from std_msgs.msg import Header

CV_ENCODING_TO_DTYPE_MAP = {"8U": np.uint8, "8S": np.int8,    "16U":np.uint16, "16S":np.int16, 
                            "32S":np.int32, "32F":np.float32, "64F":np.float64}
DTYPE_TO_CV_ENCODING_MAP = {np.uint8:"8U",  np.int8:"8S",     np.uint16:"16U", np.int16:"16S",
                            np.int32:"32S", np.float32:"32F", np.float64:"64F"}

def _encoding_to_dtype_with_channels(encoding):
    if(encoding.find('bayer') != -1 or encoding == 'yuv422'):
        raise "Sorry, I don't know this encoding!"

    channels = 0
    dtype = np.uint8
    if(encoding.find('C') == -1):
        if(encoding.find('rgba')!=-1 or encoding.find('bgra')!=-1):
            channels = 4
        elif(encoding.find('mono')!=-1):
            channels = 1
        else:
            channels = 3
        if(encoding.find('8')!=-1):
            dtype = np.uint8
        else:
            dtype = np.uint16
    else:
        dtype = CV_ENCODING_TO_DTYPE_MAP[encoding.split('C')[0]]
        channels = int(encoding[-1])
    
    return dtype, channels

def convert_Image_to_nparray(img_msg, to_rgb = True):
    dtype, n_channels = _encoding_to_dtype_with_channels(img_msg.encoding)
    if n_channels == 1:
        im = np.ndarray(shape=(img_msg.height, img_msg.width), dtype=dtype, buffer=img_msg.data)
    else:
        im = np.ndarray(shape=(img_msg.height, img_msg.width, n_channels), dtype=dtype, buffer=img_msg.data)

    if to_rgb and img_msg.encoding.find('bgr') != -1:
        if(n_channels == 3):
            im = im[...,::-1]
        elif(n_channels == 4):
            im = im[...,(2,1,3,0)]

    return im

def convert_nparray_to_Image(frame):
    """
    Input:
        frame: nparrays, cv2, bgr8
    Output:
        ros_frame: sensor_msgs.msg
    """
    ros_frame = Image()
    header = Header(stamp = rospy.Time.now())
    header.frame_id = "Camera"
    ros_frame.header = header
    ros_frame.width = frame.shape[1]    # 640
    ros_frame.height = frame.shape[0]   # 480
    ros_frame.encoding = "rgb8"
    ros_frame.step = 1920
    ros_frame.data = frame.tostring()
    return ros_frame